#export LD_PRELOAD=/usr/lib/libmpi.so python
from libatomism import *
from plotPES import *   
from repro import *   
import numpy as np
import scipy.stats as stats

msLogger.setPriorityMin(Priority.ERROR)

unit=msUnitsManager.New("Angstrom Degree amu kcal/mol ps")

zmat=msZMat.New(unit).setId("zmat").set("C O 0 1.15 O 0 5. 1 120 H 2 0.98 0 90 1 180.")
system=msSystem.New(unit).addEntity(zmat)

#viewer=msVisualizerAtoms.New()
#viewer.watch(system)

gcoors = msGeneralizedCoordinates.New(unit)
q0 = msScalarVariable.New("Angstrom").set(2., 1.45, 2.5, 0.05, 0.05).setId("q0")
q1 = msScalarVariable.New("Degree").set(180, 0, 181, 10, 10).setId("q1")
gcoors.addVar(q0).addVar(q1)

kinfct0=msParser.New("Angstrom").setCoordinates(gcoors).setExpression("q0").setId("f=q0")
kinfct1=msParser.New("Degree").setCoordinates(gcoors).setExpression("q1").setId("f=q1")

kinop=msKineticOperator.New(unit).set(gcoors,system).addDynamicDof(zmat,1,kinfct0).addDynamicDof(zmat,5,kinfct1)

epot = pesHOCO.New().setCoordinates(gcoors)
epot.setKineticOperator(kinop)
repro = setrepro(kinop,5,15)

database = msDiracs.New("kcal/mol").setCoordinates(gcoors)

with open('DataAI.txt') as f:
    w = f.readline().split() 	# read first line
    array = []
    for line in f: 		# read rest of lines
        array = [float(x) for x in line.split()]
	value = (array[len(array)-1] + 188.656188) * 627.51
	array.remove(array[len(array)-1])
	vec = VectorOfDouble()
        vec.extend( arg for arg in array)
        database.addDirac(vec, value )

invpb=uqBayesianScalarModel.New()
invpb.setModel( epot )
invpb.setDatabase( database )

invpb.setParameter("a0",0,10)
invpb.setParameter("a1",0,20)
invpb.setParameter("a2",10,20)
invpb.setParameter("a3",0,2)
invpb.setParameter("a4",0,60)
invpb.setParameter("a5",0,2)
invpb.setParameter("a6",0,20)

msLogger.setPriorityMin(Priority.INFO)
msLogger.setVerboseLevel(2)
NSample = 10000
invpb.getParameters().setParameter("ChainLength",NSample,"")
invpb.getParameters().setParameter("ChainLength",NSample,"")
invpb.getParameters().setParameter("DisplayStepRatio",0.5,"")

invpb.infer()

postPDFsample = invpb.getPosteriorPDF()
param 	      = postPDFsample.getUqParameters()
postPDFsample.set95CI()
print param

fig = plt.figure()
#fig, axes = plt.subplots(2, 2)

plotPES3d(kinop,repro, fig.add_subplot(2, 2, 1,projection='3d') )  

histoCorrel(postPDFsample,0,1,fig.add_subplot(2, 2, 2))

for i in range(0,50):
    postPDFsample.getRndPoint()
    plotPES(epot,fig.add_subplot(2, 2, 3))  
 
plotRepro(kinop,repro,fig.add_subplot(2, 2, 3))


msLogger.setPriorityMin(Priority.ERROR)
motion = msLagrangian.New(unit).setEqOfMotion(kinop,epot)

s = 1.8
q0.set(s, 1.45, 2.5, 0.0, 0.0).setId("q0")
q1.set(0, 0, 180, 1., 1.).setId("q1")
postPDFsample.getRndPoint()

ax=fig.add_subplot(2, 2, 4)
for i in range(0,10):
   postPDFsample.getRndPoint() 
   emin = epot.evaluate(s,180)
   T=[]
   Q=[]
   for j in range(0,10):
       t = 100 + 300*j
       T.append(t)
       q= motion.Q(t)*exp(emin/(1.98e-3*t))
       Q.append( q )
       print t,q
   ax.plot(T,Q,'-')

ax.axis((0,3000,0,6))
ax.set_xlabel("Temperature [K]",fontsize=20)
ax.set_ylabel("Partition function",fontsize=20)
ax.tick_params(labelsize=20)
fig.set_size_inches(20,20)
fig.savefig('HOCO_0.png')
fig.show()


